from ftplib import FTP
from tqdm import tqdm
import os

def ftp_connect():
    try:
        host = input("输入FTP服务器主机名或IP地址: ")
        username = input("输入用户名: ")
        password = input("输入密码: ")

        ftp = FTP(host)
        ftp.login(username, password)
        return ftp
    except Exception as e:
        print(f"连接到FTP服务器失败: {e}")
        return None

def list_files(ftp):
    try:
        files = ftp.nlst()
        print("\n文件列表:")
        for file in files:
            print(file)
    except Exception as e:
        print(f"列出文件失败: {e}")

def download_file(ftp, remote_path, local_path):
    try:
        with open(local_path, 'wb') as local_file:
            ftp.sendcmd('TYPE I')  # 切换到二进制传输模式
            total_size = ftp.size(remote_path)
            downloaded = 0
            progress = tqdm(total=total_size, unit="B", unit_scale=True)

            def callback(data):
                nonlocal downloaded
                downloaded += len(data)
                local_file.write(data)
                progress.update(len(data))
                tqdm.write(f"Downloading... ({downloaded}/{total_size} bytes)")

            ftp.retrbinary('RETR ' + remote_path, callback)
            progress.close()
        print(f"下载文件 {remote_path} 成功")
    except Exception as e:
        print(f"下载文件失败: {e}")

def upload_file(ftp, local_path, remote_path):
    try:
        with open(local_path, 'rb') as local_file:
            total_size = os.path.getsize(local_path)
            uploaded = 0
            progress = tqdm(total=total_size, unit="B", unit_scale=True)

            def callback(data):
                nonlocal uploaded
                uploaded += len(data)
                ftp.storbinary('STOR ' + remote_path, data)
                progress.update(len(data))
                tqdm.write(f"Uploading... ({uploaded}/{total_size} bytes)")

            ftp.storbinary('STOR ' + remote_path, local_file, callback)
            progress.close()
        print(f"上传文件 {local_path} 成功")
    except Exception as e:
        print(f"上传文件失败: {e}")

def main():
    ftp = ftp_connect()
    if ftp:
        while True:
            print("\n选择操作:")
            print("1. 列出文件")
            print("2. 下载文件")
            print("3. 上传文件")
            print("4. 退出")
            choice = input("输入数字以执行操作: ")
            
            if choice == '1':
                list_files(ftp)
            elif choice == '2':
                remote_path = input("输入远程文件路径: ")
                local_path = input("输入本地文件路径: ")
                download_file(ftp, remote_path, local_path)
            elif choice == '3':
                local_path = input("输入本地文件路径: ")
                remote_path = input("输入远程文件路径: ")
                upload_file(ftp, local_path, remote_path)
            elif choice == '4':
                ftp.quit()
                print("已退出FTP客户端")
                break
            else:
                print("无效的选择，请重新输入。")

if __name__ == "__main__":
    main()
